ARG IMAGE_NAME
FROM ${IMAGE_NAME}:{{ cuda.version.release_label }}-base-{{ cuda.os.distro }}{{ cuda.os.version }}{{ cuda.image_tag_suffix if "image_tag_suffix" in cuda }} as base

ENV NV_CUDA_LIB_VERSION {{ cuda.version.release_label }}-1

{% if "x86_64" in cuda %}
FROM base as base-amd64

ENV NV_NVTX_VERSION {{ cuda.x86_64.components.nvtx.version }}
{% if "libnpp" in cuda.x86_64.components %}
    {% set has_libnpp_package = true %}
ENV NV_LIBNPP_VERSION {{ cuda.x86_64.components.libnpp.version }}
ENV NV_LIBNPP_PACKAGE libnpp-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBNPP_VERSION}
{% endif %}
ENV NV_LIBCUBLAS_VERSION {{ cuda.x86_64.components.libcublas.version }}
{% if "libnccl2" in cuda.x86_64.components %}
    {% set has_libnccl_package = true %}
ENV NV_LIBNCCL_PACKAGE_NAME libnccl
ENV NV_LIBNCCL_PACKAGE_VERSION {{ cuda.x86_64.components.libnccl2.version }}
ENV NV_LIBNCCL_VERSION {{ cuda.x86_64.components.libnccl2.version[:-2] }}
ENV NCCL_VERSION {{ cuda.x86_64.components.libnccl2.version[:-2] }}
ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}-${NV_LIBNCCL_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}
{% endif -%}

{% endif %}
{% if "arm64" in cuda %}
FROM base as base-arm64

ENV NV_NVTX_VERSION {{ cuda.arm64.components.nvtx.version }}
{% if "libnpp" in cuda.arm64.components %}
    {% set has_libnpp_package = true %}
ENV NV_LIBNPP_VERSION {{ cuda.arm64.components.libnpp.version }}
ENV NV_LIBNPP_PACKAGE libnpp-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBNPP_VERSION}
{% endif %}
ENV NV_LIBCUBLAS_VERSION {{ cuda.arm64.components.libcublas.version }}
{% if "libnccl2" in cuda.arm64.components %}
    {% set has_libnccl_package = true %}
ENV NV_LIBNCCL_PACKAGE_NAME libnccl
ENV NV_LIBNCCL_PACKAGE_VERSION {{ cuda.arm64.components.libnccl2.version }}
ENV NV_LIBNCCL_VERSION {{ cuda.arm64.components.libnccl2.version[:-2] }}
ENV NCCL_VERSION {{ cuda.arm64.components.libnccl2.version[:-2] }}
ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}-${NV_LIBNCCL_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}
{% endif -%}

{% endif %}
{% if "ppc64le" in cuda %}
FROM base as base-ppc64le

ENV NV_NVTX_VERSION {{ cuda.ppc64le.components.nvtx.version }}
{% if "libnpp" in cuda.ppc64le.components %}
    {% set has_libnpp_package = true %}
ENV NV_LIBNPP_VERSION {{ cuda.ppc64le.components.libnpp.version }}
ENV NV_LIBNPP_PACKAGE libnpp-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBNPP_VERSION}
{% endif %}
ENV NV_LIBCUBLAS_VERSION {{ cuda.ppc64le.components.libcublas.version }}
{% if "libnccl2" in cuda.ppc64le.components %}
    {% set has_libnccl_package = true %}
ENV NV_LIBNCCL_PACKAGE_NAME libnccl
ENV NV_LIBNCCL_PACKAGE_VERSION {{ cuda.ppc64le.components.libnccl2.version }}
ENV NV_LIBNCCL_VERSION {{ cuda.ppc64le.components.libnccl2.version[:-2] }}
ENV NCCL_VERSION {{ cuda.ppc64le.components.libnccl2.version[:-2] }}
ENV NV_LIBNCCL_PACKAGE ${NV_LIBNCCL_PACKAGE_NAME}-${NV_LIBNCCL_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}

{% endif %}
{% endif -%}

FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <sw-cuda-installer@nvidia.com>"

RUN yum install -y \
    cuda-libraries-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_CUDA_LIB_VERSION} \
    cuda-nvtx-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_NVTX_VERSION} \
    {% if has_libnpp_package %}
    ${NV_LIBNPP_PACKAGE} \
    {% endif %}
    libcublas-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBCUBLAS_VERSION} \
    {% if has_libnccl_package %}
    ${NV_LIBNCCL_PACKAGE} \
    {% endif %}
    && yum clean all \
    && rm -rf /var/cache/yum/*

# Add entrypoint items
COPY entrypoint.d/ /opt/nvidia/entrypoint.d/
COPY nvidia_entrypoint.sh /opt/nvidia/
ENV NVIDIA_PRODUCT_NAME="CUDA"
{% if cuda.end_of_life %}
ENV NVIDIA_CUDA_END_OF_LIFE=1
{% endif %}
ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]
